# -*- coding: utf-8 -*-
# @Function: 数据导入器
# @Description: 负责将Excel文件中的问答对导入到数据库，包括文件读取、分词处理和数据库插入
# @Usage: 被all.py导入使用，执行数据导入流程
# @Dependencies: Z_config.py, db_manager.py, excel_reader.py, tokenizer.py

import os
import csv
import json
from typing import List, Tuple
from Z_config import RESERVOIR_MAPPING
from Z_db_manager import DBManager
from Z_excel_reader import ExcelReader
from Z_tokenizer import Tokenizer

class QAImporter:
    def __init__(self):
        """初始化导入器"""
        self.db_manager = DBManager()
        self.excel_reader = ExcelReader()
        self.tokenizer = Tokenizer()

    def process_csv_tables(self, csv_file_path: str, category_id: str) -> bool:
        """处理CSV文件中的表格数据，提取第一行第二个单元格有数据的表格

        Args:
            csv_file_path (str): CSV文件路径
            category_id (str): 类别ID

        Returns:
            bool: 是否成功
        """
        try:
            print(f"\n处理CSV表格文件: {os.path.basename(csv_file_path)}")
            
            # 读取CSV文件
            with open(csv_file_path, 'r', encoding='utf-8', errors='ignore') as file:
                csv_reader = csv.reader(file)
                all_rows = list(csv_reader)
            
            if not all_rows:
                print("CSV文件为空")
                return False
            
            # 解析表格
            tables = self._parse_csv_tables(all_rows)
            
            # 过滤出第一行第二个单元格有数据的表格
            type2_tables = []
            for table in tables:
                if len(table) > 0 and len(table[0]) > 1 and table[0][1].strip():
                    # 第一行第二个单元格有数据，这是类型二表格
                    keywords = table[0][1].strip()  # 关键词来自第一行第二个单元格
                    
                    # 构建表格JSON数据（排除第一行第二个单元格）
                    table_data = []
                    for i, row in enumerate(table):
                        if i == 0:
                            # 第一行：保留第一个单元格，第二个单元格置空，其余保持不变
                            modified_row = [row[0] if len(row) > 0 else ""] + [""] + row[2:] if len(row) > 2 else [row[0] if len(row) > 0 else "", ""]
                            table_data.append(modified_row)
                        else:
                            table_data.append(row)
                    
                    table_json = {
                        "headers": table_data[0] if table_data else [],
                        "rows": table_data[1:] if len(table_data) > 1 else [],
                        "total_rows": len(table_data),
                        "total_cols": len(table_data[0]) if table_data else 0
                    }
                    
                    type2_tables.append((keywords, table_json, category_id, os.path.basename(csv_file_path)))
            
            if type2_tables:
                # 插入表格数据到数据库
                self.db_manager.insert_table_data(type2_tables)
                print(f"成功处理 {len(type2_tables)} 个类型二表格")
            else:
                print("未找到类型二表格（第一行第二个单元格有数据的表格）")
            
            return True
            
        except Exception as e:
            print(f"处理CSV表格文件失败: {e}")
            return False

    def _parse_csv_tables(self, rows: List[List[str]]) -> List[List[List[str]]]:
        """解析CSV行数据为表格列表

        Args:
            rows (List[List[str]]): CSV行数据

        Returns:
            List[List[List[str]]]: 表格列表
        """
        tables = []
        current_table = []
        
        for row in rows:
            # 检查是否为空行（所有单元格都为空或只有空白字符）
            is_empty_row = all(cell.strip() == '' for cell in row)
            
            if is_empty_row:
                # 遇到空行，结束当前表格
                if current_table:
                    tables.append(current_table)
                    current_table = []
            else:
                # 非空行，添加到当前表格
                current_table.append(row)
        
        # 处理最后一个表格
        if current_table:
            tables.append(current_table)
        
        return tables

    def process_csv_file(self, csv_file_path: str) -> bool:
        """处理单个CSV文件

        Args:
            csv_file_path (str): CSV文件路径

        Returns:
            bool: 是否成功
        """
        try:
            # 从文件名提取水库名称
            file_name = os.path.basename(csv_file_path)
            reservoir_name = file_name.replace('.csv', '')
            
            # 查找对应的category_id
            if reservoir_name in RESERVOIR_MAPPING:
                category_id = RESERVOIR_MAPPING[reservoir_name]
                print(f"处理水库CSV: {reservoir_name} (ID: {category_id})")
                
                # 处理CSV表格
                return self.process_csv_tables(csv_file_path, category_id)
            else:
                print(f"跳过文件 {file_name} - 未找到匹配的category_id")
                return False
                
        except Exception as e:
            print(f"处理CSV文件失败: {e}")
            return False

    def process_file(self, file_path: str, category_id: str) -> bool:
        """处理单个Excel文件

        Args:
            file_path (str): Excel文件路径
            category_id (str): 类别ID

        Returns:
            bool: 是否成功
        """
        try:
            print(f"\n处理文件: {os.path.basename(file_path)}")
            
            # 步骤1: 读取Excel文件
            df = self.excel_reader.read_file(file_path)
            print("步骤1: Excel文件读取完成")
            
            # 步骤2: 解析QA数据
            qa_pairs = self.excel_reader.parse_qa_data(df)
            if not qa_pairs:
                print(f"文件 {os.path.basename(file_path)} 中没有找到有效的QA数据")
                return False
            print(f"步骤2: 解析出 {len(qa_pairs)} 个问答对")
            
            # 步骤3: 对问题进行分词
            processed_qa_pairs = []
            for question, answer in qa_pairs:
                token_data = self.tokenizer.tokenize_question(question)
                processed_qa_pairs.append((question, answer, token_data))
            print("步骤3: 问题分词处理完成")
            
            # 步骤4: 插入数据库
            self.db_manager.insert_qa_data(processed_qa_pairs, category_id)
            print("步骤4: 数据库插入完成")
            
            print(f"文件 {os.path.basename(file_path)} 处理成功")
            return True
            
        except Exception as e:
            print(f"处理文件 {os.path.basename(file_path)} 失败: {e}")
            return False

    def process_folder(self, folder_path: str) -> Tuple[int, int]:
        """处理文件夹中的所有Excel文件

        Args:
            folder_path (str): 文件夹路径

        Returns:
            Tuple[int, int]: (成功处理的文件数, 总文件数)
        """
        try:
            # 连接数据库
            self.db_manager.connect()
            
            # 创建表
            self.db_manager.create_table()
            
            # 获取所有Excel文件
            excel_files = self.excel_reader.get_excel_files(folder_path)
            if not excel_files:
                return 0, 0
                
            # 处理每个文件
            success_count = 0
            for file_path in excel_files:
                try:
                    # 提取水库名称
                    reservoir_name = self.excel_reader.extract_reservoir_name(file_path)
                    
                    # 查找对应的category_id
                    if reservoir_name in RESERVOIR_MAPPING:
                        category_id = RESERVOIR_MAPPING[reservoir_name]
                        print(f"处理水库: {reservoir_name} (ID: {category_id})")
                        
                        # 处理单个文件
                        if self.process_file(file_path, category_id):
                            success_count += 1
                    else:
                        print(f"跳过文件 {os.path.basename(file_path)} - 未找到匹配的category_id")
                        
                except Exception as e:
                    print(f"处理文件 {os.path.basename(file_path)} 失败: {e}")
                    
            return success_count, len(excel_files)
            
        except Exception as e:
            print(f"批量导入过程中出错: {e}")
            return 0, len(excel_files) if 'excel_files' in locals() else 0
        finally:
            # 关闭数据库连接
            self.db_manager.close()

    def process_specific_csv(self, csv_file_path: str = r"d:\AAA\LLM_autoQA_API\A_segmenter\csv\大伙房水库调度规程.csv") -> bool:
        """处理指定的CSV文件

        Args:
            csv_file_path (str): CSV文件路径

        Returns:
            bool: 是否成功
        """
        try:
            # 连接数据库
            self.db_manager.connect()
            
            # 创建表
            self.db_manager.create_table()
            
            # 处理CSV文件
            result = self.process_csv_file(csv_file_path)
            
            return result
            
        except Exception as e:
            print(f"处理指定CSV文件失败: {e}")
            return False
        finally:
            # 关闭数据库连接
            self.db_manager.close()